import uuid
import logging
from typing import Callable, Awaitable
from fastapi import Request, Response
from starlette.middleware.base import BaseHTTPMiddleware
from app.utils.logger import logger, set_request_id

class RequestIdMiddleware(BaseHTTPMiddleware):
    """
    请求ID中间件，为每个请求生成或传递X-Request-Id
    """
    
    async def dispatch(
        self, 
        request: Request, 
        call_next: Callable[[Request], Awaitable[Response]]
    ) -> Response:
        # 获取或生成请求ID
        request_id = request.headers.get("X-Request-Id")
        if not request_id:
            request_id = str(uuid.uuid4())
        
        # 将request_id添加到request的state中，以便在应用中使用
        request.state.request_id = request_id
        
        # 设置contextvars中的request_id
        set_request_id(request_id)

        # 继续处理请求
        response = await call_next(request)
        
        # 添加X-Request-Id到响应头
        response.headers["X-Request-Id"] = request_id

        return response